﻿using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace ZeroMQCommon
{
    public class ImageAssembler
    {
        private readonly ConcurrentDictionary<string, FrameBuffer> _buffers = new();
        private readonly TimeSpan _timeout;
        private readonly System.Timers.Timer _cleanupTimer;

        public event Action<FrameBuffer, byte[]>? ImageCompleted;
        public event Action<string>? ImageTimeout;

        public ImageAssembler(TimeSpan timeout, double cleanupIntervalMs = 10000)
        {
            _timeout = timeout;
            _cleanupTimer = new System.Timers.Timer(cleanupIntervalMs);
            _cleanupTimer.Elapsed += (_, _) => CleanupExpired();
            _cleanupTimer.Start();
        }

        public void AddChunk(string imageId, int chunkIndex, int totalChunks, byte[] data)
        {
            var buffer = _buffers.GetOrAdd(imageId, _ => new FrameBuffer { ImageId = imageId, ExpectedChunks = totalChunks });

            lock (buffer)
            {
                buffer.Chunks[chunkIndex] = data;

                if (buffer.Chunks.Count == buffer.ExpectedChunks)
                {
                    var fullImage = buffer.Chunks.OrderBy(kvp => kvp.Key).SelectMany(kvp => kvp.Value).ToArray();
                    buffer.EndTime = DateTime.Now;
                    _buffers.TryRemove(imageId, out var frame);
                    ImageCompleted?.Invoke(frame, fullImage);
                }
            }
        }

        private void CleanupExpired()
        {
            var now = DateTime.UtcNow;
            foreach (var (imageId, buffer) in _buffers)
            {
                if ((now - buffer.CreatedTime) > _timeout)
                {
                    _buffers.TryRemove(imageId, out _);
                    ImageTimeout?.Invoke(imageId);
                }
            }
        }

        public void Stop() => _cleanupTimer.Stop();
    }
}
